package blsX

import (
	"crypto/sha256"
	"kzg-ipa-study/tools"

	bls12381 "github.com/kilic/bls12-381"
)

type Setup struct {
	m, n uint64
	P    *bls12381.PointG1      // 聚合公钥
	Poly map[uint64]bls12381.Fr // hash(pi, {p1, p2, p3.....})

	MK map[uint64]*bls12381.PointG2 // 成员密钥
}

func NewSetup(m, n uint64) *Setup {
	return &Setup{
		m:  m,
		n:  n,
		MK: make(map[uint64]*bls12381.PointG2),
	}
}

// CreateP 生成聚合公钥
func (s *Setup) CreateP(pubs map[uint64][48]byte) *bls12381.PointG1 {
	if uint64(len(pubs)) != s.n {
		panic("the len of public keys wrong!")
	}

	superp := make([]byte, 0, len(pubs)*48)
	for _, p := range pubs {
		superp = append(superp, p[:]...)
	}

	G1 := bls12381.NewG1()
	sumP1 := G1.New()

	poly := make(map[uint64]bls12381.Fr, len(pubs))
	hash := sha256.New()
	for i, p := range pubs {
		hash.Write(superp)
		hash.Write(p[:])
		var tmpPoly bls12381.Fr
		tmpPoly.FromBytes(hash.Sum(nil))

		poly[i] = tmpPoly
		hash.Reset()

		pi, err := G1.FromCompressed(p[:])
		if nil != err {
			panic(err)
		}

		var tmpP1 bls12381.PointG1
		G1.MulScalar(&tmpP1, pi, &tmpPoly)
		G1.Add(sumP1, sumP1, &tmpP1)
	}

	s.Poly = poly
	s.P = sumP1
	return sumP1
}

// CreatMK 生成成员密钥
func (s *Setup) CreatMK(secrets map[uint64][32]byte) {
	if uint64(len(secrets)) != s.n {
		panic("len of secret not equal of n")
	}

	bls := bls12381.NewEngine()
	G1 := bls.G1
	G2 := bls.G2

	sumFr := bls12381.NewFr()
	for i, sec := range secrets {
		var sFr bls12381.Fr
		sFr.FromBytes(sec[:])

		poly := s.Poly[i]
		sFr.Mul(&poly, &sFr)
		sumFr.Add(sumFr, &sFr)
	}

	Pbyte := G1.ToCompressed(s.P)

	for k, _ := range secrets {

		tbyte := append(Pbyte, tools.Uint64ToByte(k)...)
		p2, err := G2.HashToCurve(tbyte, domainG2)
		if nil != err {
			panic("hash to p2 faild")
		}

		tmpp2 := G2.New()
		G2.MulScalar(tmpp2, p2, sumFr)
		s.MK[k] = tmpp2
	}
}

// MultiSign 多重签名
func (s *Setup) MultiSign(sec map[uint64][32]byte, msg []byte) []byte {
	bls := bls12381.NewEngine()
	G1 := bls.G1
	G2 := bls.G2

	Pbyte := G1.ToCompressed(s.P)
	Pbyte = append(Pbyte, msg...)
	p2, err := G2.HashToCurve(Pbyte, domainG2)
	if nil != err {
		panic(err)
	}

	sumMK := G2.New()
	sumFr := bls12381.NewFr()
	for idx, sec := range sec {
		var sFr bls12381.Fr
		sFr.FromBytes(sec[:])
		sumFr.Add(sumFr, &sFr)

		mk := s.MK[idx]
		// G2.Add(tmp, tmp, mk)
		G2.Add(sumMK, sumMK, mk)
	}

	G2.MulScalar(p2, p2, sumFr)
	G2.Add(sumMK, sumMK, p2)
	return G2.ToCompressed(sumMK)
}

func (s *Setup) MultiSignCheck(pubs map[uint64][48]byte, sign [96]byte, msg []byte) bool {
	bls := bls12381.NewEngine()
	G1 := bls.G1
	G2 := bls.G2

	pbyte := G1.ToCompressed(s.P)
	// p'
	// H(P, 1)+H(P, 3)
	sump1 := G1.New()
	sump2 := G2.New()
	for idx, pp := range pubs {
		p1, err := G1.FromCompressed(pp[:])
		if nil != err {
			panic("public to G1 failed")
		}
		G1.Add(sump1, sump1, p1)

		p_i := append(pbyte, tools.Uint64ToByte(idx)...)
		h_i, err := G2.HashToCurve(p_i, domainG2)
		if nil != err {
			panic(err)
		}
		G2.Add(sump2, sump2, h_i)

	}

	// H(P,m)
	pbyte = append(pbyte, msg...)
	H, err := G2.HashToCurve(pbyte, domainG2)
	if nil != err {
		panic(err)
	}
	// e(P’, H(P, m))⋅e(P, H(P, 1)+H(P, 3))
	bls.AddPair(sump1, H)
	bls.AddPair(s.P, sump2)
	// S'
	sumS, err := G2.FromCompressed(sign[:])
	if nil != err {
		panic(err)
	}
	// e(G, S’)
	bls.AddPairInv(G1.One(), sumS)

	return bls.Check()
}
